# -*- coding: utf-8 -*-

"""
Created on 03/23/2022
main.
@author: Kang Xiatao (kangxiatao@gmail.com)
"""


from models.base.init_utils import weights_init
from configs import *
from utils.network_utils import get_network
from utils.data_utils import get_dataloader
from train_test import *


def main():
    # ===== 超参数和模型配置 =====
    config = init_config()
    logger, writer = init_logger(config)  # logger: 日志  writer: tensorboard
    logger.info('  LR: %.5f, WD: %.5f, Epochs: %d' % (config.learning_rate, config.weight_decay, config.epoch))
    config.send_mail_str += 'LR: %.5f, WD: %.5f, Epochs: %d, Batch: %d \n' % (config.learning_rate, config.weight_decay, config.epoch, config.batch_size)

    state = None
    # ===== build/load model =====
    if config.pretrained:
        state = torch.load(config.pretrained)
        model = state['net']
        config.send_mail_str += f"use pre-trained mode -> acc:{state['acc']} epoch:{state['epoch']}\n"
        config.network = state['args'].network
        config.depth = state['args'].depth
        config.dataset = state['args'].dataset
        config.batch_size = state['args'].batch_size
        config.learning_rate = state['args'].learning_rate
        config.weight_decay = state['args'].weight_decay
        config.epoch = state['args'].epoch
        print('load model finish')
        print(state['args'])
    else:
        model = get_network(config.network, config.depth, config.dataset, use_bn=config.get('use_bn', True))
    model.cuda()

    # ===== get dataloader =====
    trainloader, testloader = get_dataloader(config, 4)
    # ===== train =====
    tr_str, print_inf = train_once(model, trainloader, testloader, config, writer, logger, state, config.lr_mode, config.optim_mode)

    # ===== 实验结果发送到邮箱 =====
    config.send_mail_str += print_inf
    config.send_mail_str += tr_str
    if 'test' not in config.exp_name:
        QQmail = mail_log.MailLogs()
        QQmail.sendmail(config.send_mail_str, header=config.send_mail_head)


if __name__ == '__main__':
    main()
